import yaml
import argparse
from pathlib import Path


class ConfigManager:
    """配置管理器"""
    
    def __init__(self, config_path: str = "configs/default_config.yaml"):
        """初始化配置管理器"""
        self.config_path = config_path
        self.config = self.load_config()
        
    def load_config(self) -> dict:
        """从YAML文件加载配置"""
        with open(self.config_path, 'r', encoding='utf-8') as file:
            return yaml.safe_load(file)
            
    def save_config(self, config: dict, path: str = None) -> None:
        """保存配置到YAML文件"""
        save_path = path or self.config_path
        with open(save_path, 'w', encoding='utf-8') as file:
            yaml.dump(config, file, default_flow_style=False, allow_unicode=True)
            
    def update_config_from_args(self, args: argparse.Namespace) -> dict:
        """从命令行参数更新配置"""
        # 根据需要更新配置
        if args.batch_size is not None:
            self.config['data']['batch_size'] = args.batch_size
        if args.learning_rate is not None:
            self.config['model']['learning_rate'] = args.learning_rate
        if args.max_epochs is not None:
            self.config['training']['max_epochs'] = args.max_epochs
        if args.experiment_name is not None:
            self.config['logging']['experiment_name'] = args.experiment_name
        if args.out_channels is not None:
            self.config['model']['out_channels'] = args.out_channels
        if args.accelerator is not None:
            self.config['training']['accelerator'] = args.accelerator
        if args.devices is not None:
            self.config['training']['devices'] = args.devices
            
        return self.config
        
    @staticmethod
    def parse_args():
        """解析命令行参数"""
        parser = argparse.ArgumentParser(description="训练2D U-Net模型")
        
        # 数据参数
        parser.add_argument("--batch_size", type=int, help="批处理大小")
        parser.add_argument("--num_workers", type=int, help="数据加载器的工作进程数")
        
        # 模型参数
        parser.add_argument("--learning_rate", type=float, help="学习率")
        parser.add_argument("--out_channels", type=int, help="输出通道数")
        
        # 训练参数
        parser.add_argument("--max_epochs", type=int, help="最大训练轮数")
        parser.add_argument("--accelerator", type=str, help="训练设备类型")
        parser.add_argument("--devices", type=int, help="设备数量")
        
        # 日志参数
        parser.add_argument("--experiment_name", type=str, help="实验名称")
        parser.add_argument("--config", type=str, default="configs/default_config.yaml", 
                          help="配置文件路径")
        
        # 其他参数
        parser.add_argument("--resume_from_checkpoint", type=str, help="从检查点恢复训练")
        
        return parser.parse_args()